﻿using System;
using System.Collections.Generic;
using System.Text;
using System.IO;
using System.Security.Cryptography;
using System.Runtime.InteropServices;

namespace SecureGRF
{
    public delegate void GRFEncryptProgressDelegate(object sender, int count, int position);

    public struct GrfItem
    {
        public string name;
        public int compSize;
        public int size;
        public uint position;
    }

    public class GRFEncrypt : IComparer<GrfItem>
    {
        private readonly string GrfSign = "Master of Magic";
        private readonly byte[] GrfCode = { 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e };
        private readonly int GrfVersion = 0x0200;
        private readonly uint GrfHeaderLen = 0x2e;
        private readonly int GrfEmptyCount = 0x07;
        private readonly byte[] ivCode = { 0x63, 0xd6, 0x51, 0xcc, 0x2f, 0xd2, 0xdc, 0x3a, 0xbc, 0x6b, 0xf4, 0x77, 0x53, 0x04, 0xd6, 0x3f };
        private readonly byte[] KeyXor = { 0x78, 0xd2, 0x05, 0x48, 0x2e, 0x29, 0x02, 0x80, 0xde, 0x70, 0x83, 0x8d, 0xa1, 0x5b, 0x18, 0x8d,
                                           0x8f, 0xc1, 0x3a, 0x40, 0x7b, 0x56, 0x0c, 0xba, 0x47, 0xbc, 0x76, 0x9f, 0xb6, 0xdd, 0xab, 0xe7 };

        private string _inputFile;
        private string _outputFile;
        private byte[] _password;

        [DllImport("zlib1.dll", CallingConvention = CallingConvention.Cdecl)]
        private static extern int uncompress(byte[] dest, ref int destLen, byte[] source, int sourceLen);
        [DllImport("zlib1.dll", CallingConvention = CallingConvention.Cdecl)]
        private static extern int compress(byte[] dest, ref int destLen, byte[] source, int sourceLen);
        [DllImport("zlib1.dll", CallingConvention = CallingConvention.Cdecl)]
        private static extern int compressBound(int sourceLen);

        public event GRFEncryptProgressDelegate OnProgress;

        public GRFEncrypt(string inputFile, string outputFile, string password)
        {
            _inputFile = inputFile;
            _outputFile = outputFile;
            _password = GeneratePassword(password);
        }

        public string DoEncrypt()
        {
            try
            {
                BinaryReader rInput = new BinaryReader(new FileStream(_inputFile, FileMode.Open, FileAccess.Read, FileShare.Read));
                try
                {
                    GrfItem[] items = LoadIndex(rInput);
                    RaiseProgress(items.Length, 0);
                    BinaryWriter wOutput = new BinaryWriter(new FileStream(_outputFile, FileMode.Create, FileAccess.ReadWrite, FileShare.None));
                    try
                    {
                        OutputHeader(wOutput, items);
                        for (int i = 0; i < items.Length; i++)
                        {
                            EncryptItem(rInput, wOutput, ref items[i]);
                            RaiseProgress(items.Length, i + 1);
                        }
                        SaveIndex(wOutput, items);
                        return "Encrypt done";
                    }
                    finally
                    {
                        wOutput.Close();
                    }
                }
                finally
                {
                    rInput.Close();
                }
            }
            catch (Exception ex)
            {
                return ex.Message;
            }
        }

        private GrfItem[] LoadIndex(BinaryReader reader)
        {
            //Check Header
            if (reader.BaseStream.Length < GrfHeaderLen)
                throw new Exception("Invalid GRF");
            reader.BaseStream.Position = 0;
            byte[] bufGrfSign = Encoding.ASCII.GetBytes(GrfSign);
            if (!CompareBytes(bufGrfSign, reader.ReadBytes(bufGrfSign.Length)))
                throw new Exception("Invalid GRF");
            reader.BaseStream.Position = 0x1e;
            uint posIndex = reader.ReadUInt32() + GrfHeaderLen;
            int itemCount2 = reader.ReadInt32();
            int itemCount = reader.ReadInt32();
            itemCount = itemCount - itemCount2 - GrfEmptyCount;
            if (itemCount <= 0) throw new Exception("Empty GRF");
            if (reader.ReadInt32() != GrfVersion)
                throw new Exception("Not Version 2.00");
            //Read Index
            reader.BaseStream.Position = posIndex;
            int compSize = reader.ReadInt32();
            int size = reader.ReadInt32();
            byte[] bufComp = reader.ReadBytes(compSize);
            //Uncompress Index
            byte[] bufOrig = new byte[size];
            if (uncompress(bufOrig, ref size, bufComp, compSize) != 0)
                throw new Exception("Invalid index");
            MemoryStream msIndex = new MemoryStream(size);
            msIndex.Write(bufOrig, 0, bufOrig.Length);
            msIndex.Position = 0;
            //Load Index
            BinaryReader rIndex = new BinaryReader(msIndex);
            List<GrfItem> itemList = new List<GrfItem>(itemCount);
            for (int i = 0; i < itemCount; i++)
            {
                GrfItem item = new GrfItem();
                item.name = ReadName(rIndex);
                item.compSize = rIndex.ReadInt32();
                rIndex.ReadUInt32();
                item.size = rIndex.ReadInt32();
                rIndex.ReadByte();
                item.position = rIndex.ReadUInt32() + GrfHeaderLen;
                itemList.Add(item);
            }
            rIndex.Close();
            itemList.Sort(this);
            return itemList.ToArray();
        }

        private bool CompareBytes(byte[] a, byte[] b)
        {
            if (a.Length != b.Length) return false;
            for (int i = 0; i < a.Length; i++)
            {
                if (a[i] != b[i]) return false;
            }
            return true;
        }

        private string ReadName(BinaryReader reader)
        {
            List<byte> dataList = new List<byte>();
            while (reader.BaseStream.Position < reader.BaseStream.Length)
            {
                byte data = reader.ReadByte();
                if (data == 0) break;
                dataList.Add(data);
            }
            return Encoding.Default.GetString(dataList.ToArray());
        }

        private void RaiseProgress(int count, int position)
        {
            if (OnProgress != null) OnProgress(this, count, position);
        }

        private void OutputHeader(BinaryWriter wOutput, GrfItem[] items)
        {
            wOutput.BaseStream.Position = 0;
            wOutput.Write(Encoding.ASCII.GetBytes(GrfSign));
            wOutput.Write(GrfCode);
            wOutput.Write((uint)0);
            wOutput.Write((int)0);
            wOutput.Write(items.Length + GrfEmptyCount);
            wOutput.Write(GrfVersion);
        }

        public int Compare(GrfItem x, GrfItem y)
        {
            if (x.position < y.position) return -1;
            else if (x.position == y.position) return 0;
            else return 1;
        }

        private byte[] GeneratePassword(string password)
        {
            byte[] bufStr = Encoding.Unicode.GetBytes(password);
            SHA256 sha256 = SHA256Managed.Create();
            byte[] bufHash = sha256.ComputeHash(bufStr);
            for (int i = 0; i < bufHash.Length; i++) bufHash[i] ^= KeyXor[i];
            return bufHash;
        }

        private void EncryptItem(BinaryReader rInput, BinaryWriter wOutput, ref GrfItem item)
        {
            rInput.BaseStream.Position = item.position;
            byte[] buf = rInput.ReadBytes(item.compSize);
            MemoryStream ms = new MemoryStream(item.compSize);
            Rijndael aes = RijndaelManaged.Create();
            CryptoStream cs = new CryptoStream(ms, aes.CreateEncryptor(_password, ivCode), CryptoStreamMode.Write);
            cs.Write(buf, 0, buf.Length);
            cs.FlushFinalBlock();
            ms.Position = 0;
            buf = new byte[(int)ms.Length];
            ms.Read(buf, 0, buf.Length);
            ms.Close();
            wOutput.BaseStream.Position = wOutput.BaseStream.Length;
            item.position = (uint)wOutput.BaseStream.Position;
            item.compSize = buf.Length;
            wOutput.Write(buf);
        }

        private void SaveIndex(BinaryWriter wOutput, GrfItem[] items)
        {
            //Save Index
            MemoryStream ms = new MemoryStream();
            BinaryWriter wIndex = new BinaryWriter(ms);
            foreach (GrfItem item in items)
            {
                wIndex.Write(Encoding.Default.GetBytes(item.name));
                wIndex.Write((byte)0);
                wIndex.Write(item.compSize);
                wIndex.Write(item.compSize);
                wIndex.Write(item.size);
                wIndex.Write((byte)1);
                wIndex.Write((uint)(item.position - GrfHeaderLen));
            }
            ms.Position = 0;
            //Compress Index
            int size = (int)ms.Length;
            byte[] bufOrig = new byte[size];
            ms.Read(bufOrig, 0, bufOrig.Length);
            ms.Close();
            int compSize = compressBound(size);
            byte[] bufComp = new byte[compSize];
            compress(bufComp, ref compSize, bufOrig, size);
            //Write Index
            wOutput.BaseStream.Position = wOutput.BaseStream.Length;
            uint pos = (uint)wOutput.BaseStream.Position;
            wOutput.Write(compSize);
            wOutput.Write(size);
            wOutput.Write(bufComp, 0, compSize);
            //Write Index Position
            wOutput.BaseStream.Position = 0x1e;
            wOutput.Write((uint)(pos - GrfHeaderLen));
        }
    }
}
